# fix_ldap_gids.py
import argparse
import sys
from ldap3 import Server, Connection, ALL, MODIFY_REPLACE
from sqlalchemy import create_engine, text

# --- 1. 配置信息 (请根据您的环境核对) ---

# LDAP 服务器配置
LDAP_SERVER_URI = 'ldap://127.0.0.1:389'
LDAP_BIND_DN = "cn=admin,dc=ai,dc=ustc,dc=edu,dc=cn"
LDAP_BIND_PASSWORD = "Ldap@UstCAi"
LDAP_BASE_DN = "dc=ai,dc=ustc,dc=edu,dc=cn"
LDAP_USER_SEARCH_BASE = f"ou=People,{LDAP_BASE_DN}"
LDAP_GROUP_SEARCH_BASE = f"ou=Group,{LDAP_BASE_DN}"

# 数据库文件路径 (假设在 instance 目录下)
DATABASE_URI = 'sqlite:///instance/app.db'

def get_approved_students_from_db(engine):
    """
    从数据库获取所有已批准学生及其导师的权威映射。
    使用正确的字符串条件进行查询。
    """
    print("Step 1: 从数据库加载已批准学生与导师的正确关系...")
    # 关键修正：使用字符串 'APPROVED' 和 'STUDENT' 进行查询
    query = """
        SELECT username, supervisor_name
        FROM application
        WHERE status = 'APPROVED' AND user_type = 'STUDENT';
    """
    try:
        with engine.connect() as connection:
            result = connection.execute(text(query))
            mapping = {row[0]: row[1] for row in result}
            print(f" -> 成功加载 {len(mapping)} 条已批准的学生记录。")
            return mapping
    except Exception as e:
        print(f"[错误] 无法连接或查询数据库: {e}")
        sys.exit(1)

def get_ldap_data(conn):
    """从LDAP获取所有用户和组的信息。"""
    print("Step 2: 从LDAP加载用户和组的当前状态...")
    
    # 获取组信息
    conn.search(LDAP_GROUP_SEARCH_BASE, '(objectClass=posixGroup)', attributes=['cn', 'gidNumber'])
    groups = conn.entries
    group_name_to_gid = {g.cn.value: int(g.gidNumber.value) for g in groups}
    gid_to_group_name = {int(g.gidNumber.value): g.cn.value for g in groups}

    # 获取用户信息
    conn.search(LDAP_USER_SEARCH_BASE, '(objectClass=posixAccount)', attributes=['uid', 'gidNumber'])
    users = {u.uid.value: {'gidNumber': int(u.gidNumber.value), 'dn': u.entry_dn} for u in conn.entries}

    print(f" -> 成功加载 {len(users)} 个用户和 {len(groups)} 个组。")
    return users, group_name_to_gid, gid_to_group_name

def main(dry_run):
    """主执行函数"""
    if dry_run:
        print("\n*** 正在以【演习模式】运行，不会对LDAP进行任何实际修改。 ***\n")
    else:
        print("\n*** 警告：正在以【执行模式】运行，将直接修改LDAP服务器数据！ ***\n")

    db_engine = create_engine(DATABASE_URI)
    approved_students = get_approved_students_from_db(db_engine)

    if not approved_students:
        print("数据库中没有找到已批准的学生，无需进行校验。")
        return

    try:
        server = Server(LDAP_SERVER_URI, get_info=ALL)
        conn = Connection(server, user=LDAP_BIND_DN, password=LDAP_BIND_PASSWORD, auto_bind=True)
    except Exception as e:
        print(f"[错误] 无法连接到LDAP服务器: {e}")
        sys.exit(1)
        
    try:
        ldap_users, group_name_to_gid, gid_to_group_name = get_ldap_data(conn)
        
        print("\nStep 3: 开始校验每个已批准学生的GID...\n" + "="*50)
        
        mismatch_count = 0
        for username, supervisor_name in approved_students.items():
            print(f"\n--- 正在校验用户: {username} (导师: {supervisor_name}) ---")

            if username not in ldap_users:
                print(f"  [警告] 用户 '{username}' 在数据库中已批准，但在LDAP中不存在。跳过。")
                continue

            # 确定正确的组和GID
            correct_group_name = f"group_{supervisor_name}"
            if correct_group_name not in group_name_to_gid:
                print(f"  [错误] 对应的导师组 '{correct_group_name}' 在LDAP中不存在。请先确保导师账户已正确配置。跳过。")
                continue
            
            correct_gid = group_name_to_gid[correct_group_name]
            
            # 获取当前的GID
            current_gid = ldap_users[username]['gidNumber']
            user_dn = ldap_users[username]['dn']

            # 对比并修正
            if current_gid == correct_gid:
                print(f"  [OK] GID正确。当前: {current_gid} ({correct_group_name})")
            else:
                mismatch_count += 1
                current_group_name = gid_to_group_name.get(current_gid, f"不存在的组(GID={current_gid})")
                print(f"  [不匹配] GID错误！")
                print(f"    - 当前 GID: {current_gid} (指向 -> {current_group_name})")
                print(f"    - 正确 GID: {correct_gid} (应该指向 -> {correct_group_name})")
                
                if not dry_run:
                    print(f"    -> [正在修正] 正在将用户 '{username}' 的gidNumber更新为 {correct_gid} ...")
                    conn.modify(user_dn, {'gidNumber': [(MODIFY_REPLACE, [str(correct_gid)])]})
                    if conn.result['result'] == 0:
                        print("    => [成功] 用户主组ID已更新。")
                    else:
                        print(f"    => [失败] 更新失败: {conn.result}")

        print("\n" + "="*50 + f"\n校验完成。共发现 {mismatch_count} 处不匹配。")

    finally:
        conn.unbind()
        print("已断开LDAP连接。")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="校验并修正LDAP中已批准学生的gidNumber，使其与数据库中的导师信息保持一致。",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        '--execute',
        action='store_true',
        help="""实际执行LDAP修改操作。
如果没有此参数，脚本将以安全的“演习模式”(Dry Run)运行，只报告不匹配项。"""
    )
    args = parser.parse_args()

    main(dry_run=not args.execute)